#include "blas.h"
#include <string.h>
#include <float.h>
#include <math.h>

void softmax(float *input, int n, float temp, int stride, float *output) {
    int i;
    float sum = 0;
    float largest = -FLT_MAX;
    for (i = 0; i < n; ++i) {
        if (input[i * stride] > largest) largest = input[i * stride]; 
    }

    for (i = 0; i < n; ++i) {
        float e = exp(input[i * stride] / temp - largest / temp);
        sum += e;
        output[i * stride] = e;
    }

    for (i = 0; i < n; ++i) {
        output[i * stride] /= sum;
    }
}

void softmax_cpu(float *input, int n, int batch, int batch_offset, int groups, int group_offset, 
    int stride, float temp, float *output) {
    int g, b;
    for (b = 0; b < batch; ++b) {
        for (g = 0; g < groups; ++g) {
            softmax(input + b * batch_offset + g * group_offset, n, temp, stride, output + b * batch_offset + g * group_offset);
        }
    } 
}

void flatten(float *x, int size, int layers, int batch, int forward) {
    float *swap = calloc(size * layers * batch, sizeof(float));
    int i, c, b;
    for (b = 0; b < batch; ++b) {
        for (c = 0; c < layers; ++c) {
            for (i = 0; i < size; ++i) {
                int i1 = b * layers * size + c * size + i;
                int i2 = b * layers * size + i * layers + c;
                if (forward) swap[i2] = x[i1];
                else swap[i1] = x[i2];
            }
        }
    }
    memcpy(x, swap, size * layers * batch * sizeof(float));
    free(swap);
}

void scale_bias(float *output, float *scales, int batch, int n, int size) {
    int i, j, b;
    for (b = 0; b < batch; b++) {
        for (i = 0; i < n; i++) {
            for (j = 0; j < size; j++) {
                output[(b * n + i) * size + j] *= scales[i];
            }
        }
    }
}

void mean_cpu(float *x, int batch, int filters, int spatial, float *mean) {
    float scale = 1. / (batch * spatial);
    int i, j, k;
    for (i = 0; i < filters; i++) {
        mean[i] = 0;
        for (j = 0; j < batch; j++) {
            for (k = 0; k < spatial; k++) {
                int index = j * filters * spatial + i * spatial + k;
                mean[i] += x[index];
            }
        }
        mean[i] *= scale;
    }
}

void variance_cpu(float *x, float *mean, int batch, int filters, int spatial, float *variance) {
    float scale = 1. / (batch * spatial - 1);
    int i, j, k;
    for (i = 0; i < filters; i++) {
        variance[i] = 0;
        for (j = 0; j < batch; j++) {
            for (k = 0; k < spatial; k++) {
                int index = j * filters * spatial + i * spatial + k;
                variance[i] += pow((x[index] - mean[i]), 2);
            }
        }
        variance[i] *= scale;
    }
}

void normalize_cpu(float *x, float *mean, float *variance, int batch, int filters, int spatial) {
    int b, f, i;
    for (b = 0; b < batch; b++) {
        for (f = 0; f < filters; f++) {
            for (i = 0; i < spatial; i++) {
                int index = b * filters * spatial + f * spatial + i;
                x[index] = (x[index] - mean[f]) / (sqrt(variance[f]) + 0.000001f);
            }
        }
    }
}

void const_cpu(int N, float ALPHA, float *X, int INCX) {
    int i;
    for (i = 0; i < N; ++i) X[i * INCX] = ALPHA;
}

void mul_cpu(int N, float *X, int INCX, float *Y, int INCY) {
    int i;
    for (i = 0; i < N; ++i) Y[i * INCY] *= X[i * INCX]; 
}

void pow_cpu(int N, float ALPHA, float *X, int INCX, float *Y, int INCY) {
    int i;
    for (i = 0; i < N; ++i) Y[i * INCY] = pow(X[i * INCX], ALPHA);
}

void axpy_cpu(int N, float ALPHA, float *X, int INCX, float *Y, int INCY) {
    int i;
    for (i = 0; i < N; ++i) Y[i * INCY] += ALPHA * X[i * INCX];
}

void scal_cpu(int N, float ALPHA, float *X, int INCX) {
    int i;
    for (i = 0; i < N; ++i) X[i * INCX] *= ALPHA;
}

void fill_cpu(int N, float ALPHA, float *X, int INCX) {
    int i;
    for (i = 0; i < N; ++i) X[i * INCX] = ALPHA;
}

void copy_cpu(int N, float *X, int INCX, float *Y, int INCY) {
    int i;
    for (i = 0; i < N; ++i) Y[i * INCY] = X[i * INCX];
}

void logistic_x_ent_cpu(int n, float *pred, float *truth, float *delta, float *error) {
    int i;
    for (i = 0; i < n; ++i) {
        float t = truth[i];
        float p = pred[i];
        error[i] = -t * log(p) - (1 - t) * log(1 - p);
        delta[i] = t - p;
    }
}

void softmax_x_ent_cpu(int n, float *pred, float *truth, float *delta, float *error) {
    for (int i = 0; i < n; ++i) {
        float t = truth[i];
        float p = pred[i];
        error[i] = (t) ? -log(p) : 0;
        delta[i] = t - p;
    }
}